#include "decomp.h"
#include "cktsat/cktcover.h"
#include <iostream>
#include <tr1/unordered_set>

using std::cout;
using std::endl;

static double drand(){return (rand()%10000)/10000.0; }
inline Lit trans( const Lit x , const Lit * lut ){
   return lut[ x.var() ] ^ (x.sign() );
}
void trans( const cls_t & x, const Lit * lut , pcls_t & y)
{
   y.clear();
   y.growTo( x.size() );
   for( unsigned i = 0 ; i < x.size() ; ++ i )
      y[i] = trans( x[i] , lut );
}
// procedure function
static void ckt2cls( CktCover & ckt , vec< cls_t > & clses )
{
   ClauseSaver< vec<Lit> > cs(clses);
   ckt.init(0);
   ckt.covering();
   ckt.assignLit( 0 );
   ckt.assignCNF( cs );
}
// 
// the function of selection input partition ( not effect the correctness)
//
static double cost( InputInfo ** in , unsigned a , unsigned b , unsigned x , unsigned y )
{
   unsigned nx=0 , ny = 0 ;
   for( unsigned i = 0 ; i < x ; ++ i )
      if( in[i][b].connect ) nx +=1;
   for( unsigned i = 0 ; i < y ; ++ i )
      if( in[a][i].connect ) ny +=1;
   return 1.0*(nx*ny)/(x*y);
}
static double balance( unsigned np , unsigned npi , InputInfo ** p )
{
   // foreach part  cost = sum 1/#input
   double cost = 1;
   unsigned pi_num[ np ]; 
   memset( pi_num , 0 , sizeof( unsigned ) * np );

   for( unsigned i = 0 ; i < np; ++ i )
      for( unsigned j = 1 ; j < npi ; ++ j )
         if( p[i][j].connect ) 
            pi_num[i] += 1 ;

   std::sort( pi_num , pi_num + np );

   for( unsigned i = 0 ; i < np-1 ; ++ i )
      cost += pi_num[i+1] - pi_num[i];

   return 1.0/cost;
}
static double independent( unsigned np , unsigned npi , InputInfo ** p )
{
   // foreach input cost = sum 1/#part
   double cost = 0;
   for( unsigned j = 1 ; j < npi ; ++j )
   {
      unsigned n =1;
      for( unsigned i = 0 ; i < np ; ++ i )
         if( p[i][j].connect ) 
            n+=1;
      cost += (1.0)/n;
   }
   return cost;
}
static void report(InputInfo ** in , unsigned x , unsigned y)
{
   double ind = independent( x , y , in );
   double bal = balance( x  ,y , in ) ;
   cout << "independent" << ind << endl;
   cout << "balance    " << bal << endl;
   cout << "cost       " << ind*bal << endl;
   for( unsigned i= 0 ; i <  x ; ++ i )
   {
      for( unsigned j = 0 ; j < y ; ++ j )
         cout << ( (in[i][j].connect)? '1':'0' ); 
      cout << endl;
   }
}
static bool select( InputInfo ** in , std::pair<unsigned,unsigned> & sel , unsigned x , unsigned y )
{
   unsigned limit = 100;
   while( limit )
   {
      unsigned a = rand()%x;
      unsigned b = rand()%(y-1)+1;
      if( in[a][b].connect == false || in[a][b].freeze )continue;
      if( drand() < cost( in , a , b , x ,y ) )
      {
         sel = std::make_pair( a , b );
         return true;
      }
      limit -=1;
   }
   return false;
}
static void freeze( InputInfo ** in ,  std::pair<unsigned,unsigned> & sel , bool b )
{
   assert( sel.second != 0 );
   in[sel.first][sel.second].freeze = b ;
}
static void connect( InputInfo ** in, std::pair<unsigned,unsigned> & sel , bool b)
{
   assert( sel.second != 0 );
   in[sel.first][sel.second].connect = b;
}
static bool decompose( MinisatProof::Solver * s, unsigned x,  InputInfo ** in , unsigned a , unsigned b )
{
   // dump the control 
   pcls_t assumpt( a*(b-1)+1 );
   assumpt[0] = in[0][0].literal();

   for( unsigned i = 0 , k = 0; i < a ; ++ i )
      for( unsigned j = 1 ; j < b ; ++ j , ++k )
         assumpt[k] = in[i][j].literal();
   
   for( unsigned i = 0 ; i < x ; ++ i )
      if( s[i].solve( assumpt ) == true ) return false;

   return true;
}
static void add_constraint( MinisatProof::Solver & s , Lit * lits , vec< cls_t > & clses )
{
   pcls_t tmp;
   for( unsigned i = 0 ; i < clses.size() ; ++ i )
   {
      trans( clses[i] , lits , tmp );
      s.addClause( tmp );
   }
}
static void addEqCls( MinisatProof::Solver & s , Lit ctrl , Lit x , Lit y )
{
   pcls_t v(3);
   v[0] = ctrl , v[1] = x , v[2] = ~y;
   s.addClause( v );
   v[1] = ~v[1] , v[2] = ~v[2];
   s.addClause( v );
}
static void addEqCtrlCls( MinisatProof::Solver & s , Lit ctrl , Lit x , Lit y , Lit o )
{
   pcls_t v(4);
   v[0] = ctrl , v[1] = x , v[2] = ~y , v[3] = ~o;
   s.addClause( v );
   v[1] = ~v[1] , v[2] = ~v[2];
   s.addClause( v );
}
void get_connect_pair( vector<std::pair<unsigned,unsigned> > & p , Bind & x , unsigned var )
{
   p.clear();
   std::tr1::unordered_set<unsigned> set;
   for( unsigned i = 0 ; i < x.size() ; ++ i )
   {
      unsigned idx = x[i][var];
      if( set.find( idx ) != set.end() ) continue;
      set.insert( idx );
      for( unsigned j = i+1 ; j < x.size() ; ++ j )
         if( x[j][var] == idx )
            p.push_back( std::make_pair( i , j ) );
   }
};
static void dump_constraint( UnsatBind & x , MinisatProof::Solver & s , InputInfo ** in , unsigned nvar , vec<cls_t> & clses , CktCover & ckt)
{
   for( unsigned i = 0 ; i < nvar ; ++ i ) s.newVar();
   unsigned fvar = x[0][0].size();
   unsigned andFlag = ckt._andFlag;

   Lit *** lits = new Lit ** [x.size()];
   
   for( unsigned i = 0 ; i < x.size () ; ++ i )
   {
      lits[i] = new Lit * [ x[i].size() ];
      for( unsigned j = 0 ; j < x[i].size() ; ++ j)
      {
         lits[i][j] = new Lit[ ckt.nVars() ];
         for( unsigned k = 0 ; k < ckt.nVars() ; ++k )
            s.newVar(),
            lits[i][j][k] = Lit( nvar ++ , 0 );
         add_constraint( s , lits[i][j] , clses );
      }
      // connect literal and in
      for( unsigned j = 0 ; j < fvar ; ++ j )
      {
         vector< std::pair< unsigned , unsigned > > connpairs;
         get_connect_pair( connpairs , x[i] , j );
         for( unsigned u = 0 ; u < connpairs.size() ; ++ u )
         {
            unsigned a = connpairs[u].first , b = connpairs[u].second; 
            for( unsigned v = 0 ; v < andFlag ; ++ v )
            {   
               Lit l =  ckt.getLit( toAigLit(v,0) );
               addEqCls( s , in[j][v].ctrl , trans( l , lits[i][a])  , trans( l , lits[i][b] ) );
            }
         }
      }
      Lit o = Lit(0,0); 
      // connect eqpair equivalent
      if( x[i].eqpair.first != x[i].eqpair.second )
      {
         o = Lit( nvar ++ , 0 ) , s.newVar();
         unsigned a = x[i].eqpair.first , b = x[i].eqpair.second;
         for( unsigned v = 0 ; v < andFlag ; ++ v )
         {
            Lit l = ckt.getLit( toAigLit(v,0) );
            addEqCtrlCls( s , in[x.pivot][v].ctrl ,trans(l,lits[i][a]),trans(l,lits[i][b]),o);
         }
      }
      pcls_t ocls;
      Lit olits[x[i].size()];
      for( unsigned j = 0 ; j < x[i].size() ; ++ j )
         olits[j] = trans( ckt.getLit( ckt._po[0] ) , lits[i][j] );
      // insert output clause 
      for( unsigned j = 0 ; j < x[i].out.size() ; ++ j )
      {
         trans( x[i].out[j] , olits , ocls );
         if( o.var() != 0 ) ocls.push( o );
         s.addClause( ocls );
      }
   }
   // connect common
   for( unsigned i = 0 ; i < x.common.size() ; ++ i )
   {
      assert( x.common[i].size() > 1 );
      Lit * ref = lits[ x.common[i][0].first ][ x.common[i][0].second ];
      for( unsigned j = 1 ; j < x.common[i].size() ; ++ j )
      {
         Lit * next = lits[x.common[i][j].first][x.common[i][j].second];
         for( unsigned k = 0 ; k < andFlag ; ++ k )
         {
            Lit l = ckt.getLit( toAigLit(k,0) );
            addEqCls( s , in[x.pivot][k].ctrl , trans( l , ref ) , trans( l , next) );
         }
      }
   }
}
static void dump_constraint( vector<UnsatBind> & x , MinisatProof::Solver * s , InputInfo ** in , unsigned nvar , CktCover & ckt )
{
   vec<cls_t> clses;
   ckt2cls( ckt , clses );
   
   for( unsigned i = 0 ; i < x.size() ; ++ i )
      dump_constraint( x[i] , s[i] , in , nvar , clses , ckt);

}
static void connect_selection( bool ** conn , CktCover & st , vector<UnsatBind> & x )
{
   // setting 
   unsigned fvar    = x[0][0].size();
   unsigned andFlag = st._andFlag;
   InputInfo ** in = allocate( fvar , andFlag , InputInfo() );

   // allocate the in variable
   unsigned nvar = 1; 
   for( unsigned i = 0 ; i < fvar ; ++ i )
   {
      in[i][0].ctrl = Lit( 0 , 0 );
      for( unsigned j = 1 ; j < andFlag ; ++ j )
         in[i][j].ctrl = Lit( nvar ++ , 0 );
   }
    
   // dump the constraint 
   MinisatProof::Solver * s = new MinisatProof::Solver[ x.size() ];  
   dump_constraint( x , s , in , nvar , st );
   
   // incremental sat
   unsigned limit = 100;
   std::pair<unsigned,unsigned> sel;
   while( limit )
   {
      if( select( in , sel , fvar , andFlag ) == false )break; 
      connect( in , sel , false );
      if( decompose( s , x.size() , in , fvar , andFlag ) == false )
      {
         connect( in , sel , true );
         freeze( in , sel , true );
      }
   };

   // return 
   conn= allocate( fvar , andFlag , true ); 
   for( unsigned i = 0 ; i < fvar ; ++ i ) 
      for( unsigned j = 1 ; j < andFlag ; ++ j ) // the connect of 0 is setted @ allocate
         conn[i][j] = in[i][j].connect;
  
   // recycle  
   delete [] s;
   for( unsigned i = 0 ; i < fvar ; ++ i )delete [] in[i]; delete [] in;
}
